#include "kvm/kvm.h"
#include "kvm/util.h"
#include "kvm/8250-serial.h"
#include "kvm/virtio-console.h"
#include "kvm/fdt.h"
#include "iie/iie-kvmtool.h"

#include <linux/kernel.h>
#include <linux/kvm.h>
#include <linux/sizes.h>

#include <linux/kvm.h>

struct kvm_ext kvm_req_ext[] = {
	{ DEFINE_KVM_EXT(KVM_CAP_ONE_REG) },
	{ 0, 0 },
};

u64 kvm__arch_default_ram_address(void)
{
	return RISCV_RAM;
}

void kvm__arch_validate_cfg(struct kvm *kvm)
{
}

bool kvm__arch_cpu_supports_vm(void)
{
	/* The KVM capability check is enough. */
	return true;
}

void kvm__init_ram(struct kvm *kvm)
{
	int err;
	u64 phys_start, phys_size;
	void *host_mem;

	phys_start	= RISCV_RAM;
	phys_size	= kvm->ram_size;
	host_mem	= kvm->ram_start;

	err = kvm__register_ram(kvm, phys_start, phys_size, host_mem);
	if (err)
		die("Failed to register %lld bytes of memory at physical "
		    "address 0x%llx [err %d]", phys_size, phys_start, err);

	kvm->arch.memory_guest_start = phys_start;
}

void kvm__arch_delete_ram(struct kvm *kvm)
{
	munmap(kvm->arch.ram_alloc_start, kvm->arch.ram_alloc_size);
}

void kvm__arch_read_term(struct kvm *kvm)
{
	serial8250__update_consoles(kvm);
	virtio_console__inject_interrupt(kvm);
}

void kvm__arch_set_cmdline(char *cmdline, bool video)
{
}

void kvm__arch_init(struct kvm *kvm)
{
	/*
	 * Allocate guest memory. We must align our buffer to 64K to
	 * correlate with the maximum guest page size for virtio-mmio.
	 * If using THP, then our minimal alignment becomes 2M.
	 * 2M trumps 64K, so let's go with that.
	 */
	kvm->ram_size = min(kvm->cfg.ram_size, (u64)RISCV_MAX_MEMORY(kvm));
	kvm->arch.ram_alloc_size = kvm->ram_size + SZ_2M;
	kvm->arch.ram_alloc_start = mmap_anon_or_hugetlbfs(kvm,
						kvm->cfg.hugetlbfs_path,
						kvm->arch.ram_alloc_size);

	if (kvm->arch.ram_alloc_start == MAP_FAILED)
		die("Failed to map %lld bytes for guest memory (%d)",
		    kvm->arch.ram_alloc_size, errno);

	kvm->ram_start = (void *)ALIGN((unsigned long)kvm->arch.ram_alloc_start,
					SZ_2M);

	madvise(kvm->arch.ram_alloc_start, kvm->arch.ram_alloc_size,
		MADV_MERGEABLE);

	madvise(kvm->arch.ram_alloc_start, kvm->arch.ram_alloc_size,
		MADV_HUGEPAGE);
}

#define FDT_ALIGN	SZ_4M
#define INITRD_ALIGN	8

//add load file struct
// struct load_file{
// 	unsigned long src_hva;
// 	unsigned long des_gpa;
// 	unsigned long file_size;
// };

bool kvm__arch_load_kernel_image(struct kvm *kvm, int fd_kernel, int fd_initrd,
				 const char *kernel_cmdline)
{
	void *pos, *kernel_end, *limit;
	unsigned long guest_addr, kernel_offset;
	ssize_t file_size;

	/* 
	 * Linux requires the initrd and dtb to be mapped inside lowmem,
	 * so we can't just place them at the top of memory.
	 */
	limit = kvm->ram_start + min(kvm->ram_size, (u64)SZ_256M) - 1;

#if __riscv_xlen == 64
	/* Linux expects to be booted at 2M boundary for RV64 */
	kernel_offset = 0x200000;
#else
	/* Linux expects to be booted at 4M boundary for RV32 */
	kernel_offset = 0x400000;
#endif
	//pos is HVA
	pos = kvm->ram_start + kernel_offset;
	kvm->arch.kern_guest_start = host_to_guest_flat(kvm, pos);
	//original code
	if(!kvm->cfg.cmode)
	{
		file_size = read_file(fd_kernel, pos, limit - pos);
		if (file_size < 0) {
			if (errno == ENOMEM)
				die("kernel image too big to fit in guest memory.");
			die_perror("kernel read");
		}
		kernel_end = pos + file_size;
	}
	else
	{
		//we use mmap to map kernel into memory
		struct stat sb;
		if (fstat(fd_kernel, &sb) == -1) {
			die("kernel image fstat not find.");
		}
		file_size = sb.st_size;
		unsigned long *kernel_addr;
		kernel_addr = (unsigned long *)mmap(NULL, sb.st_size, PROT_READ, MAP_PRIVATE, fd_kernel, 0); 
		if (kernel_addr == MAP_FAILED) {
			die("kernel image mmap return error.");
		}

		int ret;
		struct load_file load_kernel;
		load_kernel.src_hva = (unsigned long)kernel_addr;
		load_kernel.des_gpa = kvm->arch.kern_guest_start;
		load_kernel.file_size = sb.st_size;
		ret = ioctl(kvm->vm_fd, KVM_LOAD_FILE, &load_kernel);
		if(ret != 0){
			die("kernel image too big to fit in guest memory.");
		}
		kernel_end = pos + file_size;

		//continue
		pr_debug("Loaded kernel to 0x%llx (%zd bytes)  - 0x%llx",
			kvm->arch.kern_guest_start, file_size,
			host_to_guest_flat(kvm,kernel_end));
	}

	/* Place FDT just after kernel at FDT_ALIGN address */
	pos = kernel_end + FDT_ALIGN;
	guest_addr = ALIGN(host_to_guest_flat(kvm, pos), FDT_ALIGN);
	pos = guest_flat_to_host(kvm, guest_addr);
	if (pos < kernel_end)
		die("fdt overlaps with kernel image.");

	kvm->arch.dtb_guest_start = guest_addr;
	pr_debug("Placing fdt at 0x%llx - 0x%llx",
		 kvm->arch.dtb_guest_start,
		 host_to_guest_flat(kvm, limit));

	
	/* ... and finally the initrd, if we have one. */
	if (fd_initrd != -1) {
		printf("----------------------begin load initrd./!-------------------------\n");
		struct stat sb;
		unsigned long initrd_start;

		if (fstat(fd_initrd, &sb))
			die_perror("fstat");

		pos = limit - (sb.st_size + INITRD_ALIGN);
		guest_addr = ALIGN(host_to_guest_flat(kvm, pos), INITRD_ALIGN);
		pos = guest_flat_to_host(kvm, guest_addr);
		if (pos < kernel_end)
			die("initrd overlaps with kernel image.");
		
		initrd_start = guest_addr;
		file_size = sb.st_size;
		kvm->arch.initrd_guest_start = initrd_start;
		kvm->arch.initrd_size = file_size;
		unsigned long *initrd_addr;
    	initrd_addr = (unsigned long *)mmap(NULL, sb.st_size, PROT_READ, MAP_PRIVATE, fd_initrd, 0); 
    	if (initrd_addr == MAP_FAILED) {
        	die("kernel image mmap return error.");

    	}
		unsigned long pages_num = sb.st_size/PAGE_SIZE;
		printf("initrd is 0x%ld bytes!\n", file_size);
		printf("initrd need 0x%ld pages!\n", pages_num);
    	printf("initrd is mapped at address: %p\n", (void *)initrd_addr);
		struct load_file load_initrd;
		load_initrd.src_hva = (unsigned long)pos;
		load_initrd.des_gpa = initrd_start;
		load_initrd.file_size = sb.st_size;
		ioctl(kvm->vm_fd, KVM_LOAD_FILE, load_initrd);
		printf("----------------------begin load initrd!-------------------------\n");
		pr_debug("Loaded initrd to 0x%llx (%llu bytes)",
			 kvm->arch.initrd_guest_start,
			 kvm->arch.initrd_size);
	} else {
		kvm->arch.initrd_size = 0;
	}
	return true;
}

bool kvm__load_firmware(struct kvm *kvm, const char *firmware_filename)
{
	/* TODO: Firmware loading to be supported later. */
	return false;
}

int kvm__arch_setup_firmware(struct kvm *kvm)
{
	return 0;
}
